# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/27 11:37 
@Author  : ZhangShenao 
@File    : 2.使用LCEL表达式.py 
@Desc    : 使用LCEL表达式
"""

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_mistralai import ChatMistralAI
from langchain_openai import ChatOpenAI

prompt = ChatPromptTemplate.from_template('Please answer questions about {query}')

gpt_4o_mini = ChatOpenAI(model='gpt-4o-mini')  # 使用OpenAI的gpt-4o-mini模型
mixtral_8x7b = ChatMistralAI(model='Mixtral 8x7B')  # 使用MistralAI的Mixtral 8x7B模型

output_parser = StrOutputParser()

# 基于相同的流程,替换不同的LLM组件,编排成不同的Chain
gpt_4o_chain = prompt | gpt_4o_mini | output_parser
mixtral_8x7b_chain = prompt | mixtral_8x7b | output_parser

# 执行Chain,返回结果
gpt_4o_mini_content = gpt_4o_chain.invoke({'query': 'Who are you?'})
mixtral_8x7b_content = mixtral_8x7b_chain.invoke({'query': 'Who are you?'})

# 对比不同模型的结果
print(f'gpt-4o-mini: {gpt_4o_mini_content}\nmixtral_8x7b: {mixtral_8x7b_content}')
